import pymysql
import json
import os
import logging
import time
import argparse


class DatabaseComparator:
    def __init__(self, config_file):
        """Initialize the database comparator with configuration"""
        self.config = self.load_config(config_file)
        self.source_conn = None
        self.target_conn = None

        # Configure logging
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[
                logging.StreamHandler(),
                logging.FileHandler('database_comparison.log')
            ]
        )
        self.logger = logging.getLogger(__name__)

    def load_config(self, config_file):
        """Load configuration from JSON file"""
        try:
            with open(config_file, 'r') as f:
                return json.load(f)
        except Exception as e:
            self.logger.error(f"Failed to load config file: {str(e)}")
            raise

    def create_connection(self, db_config):
        """Create a database connection with retry logic"""
        max_retries = 3
        retry_delay = 2

        for attempt in range(max_retries):
            try:
                conn = pymysql.connect(
                    host=db_config['host'],
                    port=db_config['port'],
                    user=db_config['user'],
                    password=db_config['password'],
                    database=db_config['database'],
                    connect_timeout=10,
                    read_timeout=30,
                    ssl=None if not db_config.get('ssl', False) else {'ssl': {}}
                )
                return conn
            except Exception as e:
                if attempt == max_retries - 1:
                    raise
                self.logger.warning(f"Connection attempt {attempt + 1} failed: {str(e)}")
                time.sleep(retry_delay)

    def connect_databases(self):
        """Establish connections to source and target databases"""
        try:
            self.logger.info("Connecting to source database...")
            self.source_conn = self.create_connection(self.config['source_db'])

            self.logger.info("Connecting to target database...")
            self.target_conn = self.create_connection(self.config['target_db'])

            self.logger.info("Successfully connected to both databases")
        except Exception as e:
            self.logger.error(f"Database connection failed: {str(e)}")
            raise

    def get_table_counts(self, conn, table_name):
        """Get row count for a specific table with retry logic"""
        max_retries = 2
        for attempt in range(max_retries):
            try:
                with conn.cursor() as cursor:
                    cursor.execute(f"SELECT COUNT(*) FROM `{table_name}`")
                    count = cursor.fetchone()[0]
                    return count
            except pymysql.err.OperationalError as e:
                if attempt == max_retries - 1:
                    raise
                self.logger.warning(f"Retrying count for {table_name} (attempt {attempt + 1})")
                time.sleep(1)
                # Re-establish connection if needed
                if str(e).find("Packet sequence") != -1 or str(e).find("Lost connection") != -1:
                    self.logger.info("Reconnecting to database...")
                    conn.ping(reconnect=True)
            except Exception as e:
                raise

    def get_all_tables(self, conn):
        """Get list of all tables in a database"""
        try:
            with conn.cursor() as cursor:
                cursor.execute("SHOW TABLES")
                tables = [table[0] for table in cursor.fetchall()]
                return tables
        except Exception as e:
            self.logger.error(f"Failed to get table list: {str(e)}")
            raise

    def compare_table_counts(self, table_name):
        """Compare row counts for a table between source and target"""
        try:
            source_count = self.get_table_counts(self.source_conn, table_name)
            target_count = self.get_table_counts(self.target_conn, table_name)

            return {
                'table_name': table_name,
                'source_count': source_count,
                'target_count': target_count,
                'is_match': source_count == target_count,
                'difference': abs(source_count - target_count)
            }
        except Exception as e:
            self.logger.warning(f"Comparison failed for table {table_name}: {str(e)}")
            return {
                'table_name': table_name,
                'error': str(e)
            }

    def compare_all_tables(self):
        """Compare row counts for all tables between databases"""
        if not self.source_conn or not self.target_conn:
            self.connect_databases()

        try:
            source_tables = self.get_all_tables(self.source_conn)
            target_tables = self.get_all_tables(self.target_conn)

            # Identify common and unique tables
            common_tables = set(source_tables) & set(target_tables)
            source_only = set(source_tables) - set(target_tables)
            target_only = set(target_tables) - set(source_tables)

            self.logger.info(f"Common tables: {len(common_tables)}")
            self.logger.info(f"Tables only in source: {len(source_only)}")
            self.logger.info(f"Tables only in target: {len(target_only)}")

            if source_only:
                self.logger.info("Tables only in source database:")
                for table in source_only:
                    self.logger.info(f"  - {table}")

            if target_only:
                self.logger.info("Tables only in target database:")
                for table in target_only:
                    self.logger.info(f"  - {table}")

            # Compare row counts for common tables sequentially
            self.logger.info("Starting row count comparison for common tables...")

            matched_tables = []
            mismatched_tables = []
            error_tables = []

            for table_name in common_tables:
                result = self.compare_table_counts(table_name)

                if 'error' in result:
                    error_tables.append(result)
                    self.logger.warning(f"Table {result['table_name']} comparison error: {result['error']}")
                elif result['is_match']:
                    matched_tables.append(result)
                    self.logger.info(f"Table {result['table_name']}: Counts match ({result['source_count']})")
                else:
                    mismatched_tables.append(result)
                    self.logger.warning(
                        f"Table {result['table_name']}: Count mismatch "
                        f"(source: {result['source_count']}, "
                        f"target: {result['target_count']}, "
                        f"difference: {result['difference']})"
                    )

            # Summary statistics
            self.logger.info("\nComparison summary:")
            self.logger.info(f"Total tables compared: {len(common_tables)}")
            self.logger.info(f"Tables with matching counts: {len(matched_tables)}")
            self.logger.info(f"Tables with mismatched counts: {len(mismatched_tables)}")
            self.logger.info(f"Tables with comparison errors: {len(error_tables)}")

            # Save detailed results
            comparison_results = {
                'matched_tables': matched_tables,
                'mismatched_tables': mismatched_tables,
                'error_tables': error_tables,
                'source_only_tables': list(source_only),
                'target_only_tables': list(target_only)
            }
            self.save_results(comparison_results)

            return comparison_results

        except Exception as e:
            self.logger.error(f"Comparison process failed: {str(e)}")
            raise

    def save_results(self, results):
        """Save comparison results to JSON file"""
        try:
            os.makedirs('results', exist_ok=True)
            output_file = os.path.join('results', 'comparison_results.json')

            with open(output_file, 'w') as f:
                json.dump(results, f, indent=2)

            self.logger.info(f"Results saved to {output_file}")
        except Exception as e:
            self.logger.error(f"Failed to save results: {str(e)}")
            raise

    def close_connections(self):
        """Close database connections"""
        try:
            if self.source_conn:
                self.source_conn.close()
            if self.target_conn:
                self.target_conn.close()
            self.logger.info("Database connections closed")
        except Exception as e:
            self.logger.error(f"Error closing connections: {str(e)}")


def parse_args():
    """Parse command line arguments"""
    parser = argparse.ArgumentParser(description='MySQL Database Comparison Tool')
    parser.add_argument('config_file', help='Path to configuration JSON file')
    return parser.parse_args()


if __name__ == "__main__":
    try:
        args = parse_args()
        comparator = DatabaseComparator(args.config_file)
        comparator.connect_databases()
        comparison_results = comparator.compare_all_tables()
    except Exception as e:
        logging.error(f"Fatal error in main execution: {str(e)}")
    finally:
        comparator.close_connections()
